'''
数据集的读取操作
'''

import os
import sys
import torch
import torch.nn.functional as F
from torch.utils.data.dataset import Dataset
from torch.utils.data.dataloader import DataLoader
from torchvision import transforms
import numpy as np
from PIL import Image
import pandas as pds
import random

'''
自定义数据加载器
'''

class MyDataset(Dataset):
    def __init__(self,csvpath,errorcsv,traincls="train",errload=True):
        """
        csvpath: str csv文件路径  ,id,annotation,images,sig,class
        traincls: str 训练类型
        """
        errimg=[]
        if errload:
            with open(errorcsv,'r',encoding="utf-8") as fp:
                while True:
                    line=fp.readline()
                    line=line.replace(" ","").replace("\n","")
                    if len(line)<4:
                        break
                    errimg.append(line)


        csvtxt=pds.read_csv(csvpath,encoding="utf-8")# 对应的文件列表
        n=len(csvtxt)
        self.data=[] # 数据列表
        if traincls=="train": # 表示适用于训练
            for i in range(n):
                if os.path.basename(csvtxt.iloc[i]["images"]) in errimg  and errload:
                    print(os.path.basename(csvtxt.iloc[i]["images"]))
                    continue
                if csvtxt.iloc[i]["sig"]=="train":
                    self.data.append({"anpath":csvtxt.iloc[i]["annotation"],"imgpath":csvtxt.iloc[i]["images"],"class":csvtxt.iloc[i]["class"]})
        elif traincls=="vail":
                for i in range(n):
                    if os.path.basename(csvtxt.iloc[i]["images"]) in errimg and errload:
                        print(os.path.basename(csvtxt.iloc[i]["images"]))
                        continue
                    if csvtxt.iloc[i]["sig"]=="vail":
                        self.data.append({"anpath":csvtxt.iloc[i]["annotation"],"imgpath":csvtxt.iloc[i]["images"],"class":csvtxt.iloc[i]["class"]})
        elif traincls=="test":
                for i in range(n):
                    if os.path.basename(csvtxt.iloc[i]["images"]) in errimg and errload:
                        print(os.path.basename(csvtxt.iloc[i]["images"]))
                        continue
                    if csvtxt.iloc[i]["sig"]=="test":
                        self.data.append({"anpath":csvtxt.iloc[i]["annotation"],"imgpath":csvtxt.iloc[i]["images"],"class":csvtxt.iloc[i]["class"]})
        self.ImgTransfrom=transforms.Compose([
            
            transforms.ToTensor()
        ])
    def PreProcessImage(self,img):
        img=np.array(img)
        img=self.ImgTransfrom(img)
        return img
    def __getitem__(self, index):
        if not os.path.exists(self.data[index]["imgpath"]):
            index=random.randint(len(self.data))
        img=Image.open(self.data[index]["imgpath"])
        label=int(self.data[index]["class"])
        img=self.PreProcessImage(img)
        return img,label,os.path.basename(self.data[index]["imgpath"])
    def __len__(self):
        return len(self.data)


def splitCsv(csvpath):
    anncsv=pds.read_csv(csvpath,encoding="utf-8")    
    